{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "bc9bd114",
   "metadata": {},
   "source": [
    "# Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "6bf99115",
   "metadata": {},
   "outputs": [],
   "source": [
    "import jax\n",
    "from jax import numpy as jp\n",
    "from matplotlib.lines import Line2D\n",
    "from matplotlib.patches import Circle\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "try:\n",
    "  import brax\n",
    "except (ImportError, ModuleNotFoundError):\n",
    "  from IPython.display import clear_output \n",
    "  !pip install git+https://github.com/google/brax.git@main\n",
    "  clear_output()\n",
    "  import brax"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "07ed4bf1",
   "metadata": {},
   "source": [
    "# Load model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "3e0b9668",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/binh/miniforge3/envs/mujoco/lib/python3.11/site-packages/brax/io/mjcf.py:480: UserWarning: Brax System, piplines and environments are not actively being maintained. Please see MJX for a well maintained JAX-based physics engine: https://github.com/google-deepmind/mujoco/tree/main/mjx. For a host of environments that use MJX, see: https://github.com/google-deepmind/mujoco_playground.\n",
      "  warnings.warn(\n"
     ]
    }
   ],
   "source": [
    "from brax.io import mjcf\n",
    "import mujoco\n",
    "\n",
    "MODEL_PATH = '../models/lelamp/scene.xml'\n",
    "\n",
    "mj_model = mujoco.MjModel.from_xml_path(MODEL_PATH)\n",
    "mj_model.opt.solver = mujoco.mjtSolver.mjSOL_CG\n",
    "mj_model.opt.iterations = 6\n",
    "mj_model.opt.ls_iterations = 6\n",
    "\n",
    "lelamp = mjcf.load_model(mj_model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "07059a26",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[0.140913  0.0304303 0.0439213 0.040387  0.0271577 0.0121476]\n"
     ]
    }
   ],
   "source": [
    "print(lelamp.link.inertia.mass)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "f7dacbc4",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "12"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "lelamp.q_size()  # 3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d0c8ea09",
   "metadata": {},
   "outputs": [],
   "source": [
    "#@title { run: \"auto\"}\n",
    "from brax.generalized import pipeline as generalized_pipeline \n",
    "from brax.positional import pipeline as positional_pipeline \n",
    "from brax.spring import pipeline as spring_pipeline \n",
    "\n",
    "pipeline = 'generalized'  #@param [\"generalized\", \"positional\", \"spring\"]\n",
    "step_size = \"1 ms\" #@param [\".4 ms\", \"1 ms\", \"5 ms\"]\n",
    "\n",
    "substeps = {'.4 ms': 25, '1 ms': 10, '5 ms': 2}[step_size]\n",
    "pipeline = {'generalized': generalized_pipeline,\n",
    "            'positional': positional_pipeline,\n",
    "            'spring': spring_pipeline}[pipeline]\n",
    "\n",
    "lelamp = lelamp.tree_replace({'opt.timestep': 0.01 / substeps})\n",
    "\n",
    "init_q = jp.zeros(lelamp.q_size())\n",
    "state = pipeline.init(lelamp, init_q, jp.zeros(lelamp.qd_size()))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "mujoco",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
